#
#   FunctionalBohm.py
#
#   Defines objects that solve the equations of Bohmian Mechanics
#   numerically, and returns the results as functions. The code is
#   written in a functional style, mainly to use `scipy.misc.derivative'
#   (variable order, central difference, over a variable number of points).
#   We choose N=5 points, s.t. errors in first and second order are O(h^4):
#
#       df   f_{-2} - 8*f_{-1} + 8*f_{+1} - f_{+2}
#       -- = -------------------------------------
#       dx                   12*h
#

# -----------------------------------------------------------

#
#   Copyright (C) 2012
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

# -----------------------------------------------------------


# The code is primarily meant to be loaded as a module by other code.
# We only need to retrieve the object class SchrodingerNDT which contains
# all of the physics (except trajectories given an initial condition).
#
# >>> from FunctionalBohm import SchrodingerNDT
#
# but we can also run it as a main program to run a few diagnostic tests...

if __name__ == "__main__":
    print("Running Diagnostic Tests...")

# -----------------------------------------------------------

# A few requirements...

from cmath import phase as angle
from math import log
from math import sqrt
from numpy import square
from scipy.misc import derivative
from sys import float_info # to deal with infinities

# hbar = 1.0 is omitted throughout.

# ----------------------- Functionals -----------------------

# These functions have functions as input and return functions.

def phase(f): # f is a function that returns a complex number
    def wrapper(t,*xyz):
        return angle(f(t,*xyz))
    return wrapper

# Derivatives are functionals!

# We use Cartesian coordinates such that derivatives aren't too messy.

def partial(f,dx=1.0):
    # returns the derivative of f(x) as a function
    def wrapper(x,dx=dx):
        return derivative(f,x,dx, order=5) # central difference over N=5 points.
    return wrapper
    # The function g = partial(f) can take either one or two arguments: either g(x) or g(x,dx).

def laplacian1D(f,dx=1.0):
    # Again, either call laplacian(f) or laplacian(f,dx)
    # Central difference over N=5 points is O(h^3) for this `double derivative'.
    # But that's OK because our Runge-Kutta (to get trajectories) won't use this.
    def wrapper(x,dx=dx):
        return derivative(f,x,n=2,dx=dx,order=5)
    return wrapper

# Now build the N-dimensional versions:

def gradient(f,N,dx=1.0):
    # return the gradient wrt the first N dimensions of f (recursively)
    def wrapper(*args):
        if N == 1:
            if len(args) > 1:
                print("Too many arguments, ignoring all but the first")
            return [partial(f,dx)(args[0])] # only one variable
        if N > 1:
            fx = lambda x: f(x,*args[1:]) # arguments other than x curried
            fothers = lambda *others:f(args[0],*others) # x curried
            return [partial(fx,dx)(args[0])] + gradient(fothers,N-1,dx)(*args[1:])
        if N < 1:
            print("This should never happen")
            exit(0)
    return wrapper

def divergence(f,dx=1.0):
    # Returns the divergence of f.
    # f is a function of arguments *args returning some vector.
    def wrapper(*args):
        div = 0
        length = len(f(*args)) # length of the vector returned by f
        for i in range(length):
            g = lambda *args: f(*args)[i] # wrapper for the ith component, g returns a scalar
            gradg = gradient(g,length,dx) # gradient of g
            div += gradg(*args)[i] # only the ith component is nonzero (orthogonality in cartesians)
        return div
    return wrapper # wrapper is the function div(f)


def laplacian(f,N,dx=1.0):
    # return the laplacian wrt the first N dimensions of f (recursively)
    def wrapper(*args):
        if N == 1:
            if len(args) > 1:
                print("Too many arguments, ignoring all but the first")
            return laplacian1D(f,dx)(args[0]) # only one variable
        if N > 1:
            fx = lambda x: f(x,*args[1:])
            fothers = lambda *others:f(args[0],*others)
            return laplacian1D(fx,dx)(args[0]) + laplacian(fothers,N-1,dx)(*args[1:])
        if N < 1:
            print("This should never happen")
            return 0
    return wrapper

# ----------------------- Equation Solver Object -----------------------

class SchrodingerNDT: # N dimensional, time dependent
    
    # Note that the time dependence is the FIRST argument to the functions.
    # This allows us to write Psi(t,x,y,z) instead of Psi([x,y,z],t)
    # by employing the Psi(t,*x) notation in definitions.
    #    => Much more flexible
    
    def __init__(self, potential, wavefunction, mass, dimensions=3, dx=1.0, dt=1.0):
        
        self.pot = potential        # V(args) is time independent
        self.w = wavefunction       # Psi(t,*xyz)
        self.m = mass
        self.N = dimensions         # Single particle default of N=3
        self.dx = dx
        self.dt = dt
        self.__qA__ = lambda *xyz : [x*self.q for x in self.A(*xyz)] # returns q*A(x,y,z)

    def get_momentum(self):
        # The momentum in Bohmian Mechanics is
        #              dS
        # p  =  hbar * --  -  q * A(x)
        #              dx
        def wrapper(t,*args):
            S = lambda *xyz: phase(self.w)(t,*xyz) # Curry time
            p = gradient(S,N=self.N,dx=self.dx) # this funtion returns p(x).
            return p(*args)
        return wrapper
    # REEMEMBER TO RECYCLE THE qA VALUES IN THE LIST COMPREHENSION!!!
    
    def get_Imaginary_Momentum(self):
        # The weak value has an imaginary part
        #
        # Im(p) = hbar * Grad(R)/R
        #
        def wrapper(t,*xyz):
            R = lambda *args: abs(self.w(t,*args)) # R = |Psi|, Curried time
            gradR = gradient(R,self.N,self.dx)
            return [i/R(*xyz) for i in gradR(*xyz)] # hbar = 1.0
        return wrapper
    
    def get_momentum_variance(self):
        # Moyal framework gives a nonzero value to the variance around the conditional mean.
        # This means there may be deviations from the average value get_momentum().
        # The possibility of negative values is discussed in Moyal's paper.
        #
        # The momentum variance is
        #              hbar^2
        # V(p|x)  =  - ------ laplacian(ln(rho))
        #                4
        def wrapper(t,*args):
            def lnrho(*args):
                try:
                    return log(abs(self.w(t,*args))) # curry time
                except ValueError:
                    return -float_info.max # ln(0) is -infinity, for us ~ -1e308
            k = laplacian(lnrho,N=self.N,dx=self.dx) # Spatial laplacian
            return -k(*args)/2.0 # Scaling 1/4 * 2 (from ln(square) ...)
        return wrapper

    def get_Temperature(self):
        # IS THERE A MISSING MASS TERM ???
        return lambda t,*args: 2 * get_momentum_variance(t,*args) / self.N
    
    def get_Kinetic_Energy(self):
        # The Kinetic energy is
        #     p^2
        # K = ---
        #     2*m
        def wrapper(t,*args):
            p = self.get_momentum()(t,*args)
            return sum(square(p))/(2*self.m) # numpy.square accepts vector input
        return wrapper

    def get_Quantum_Potential(self):
        # The quantum potential in Bohmian Mechanics is
        #       hbar^2 laplacian(R)
        # Q = - ------ ------------
        #         2*m       R
        #
        def wrapper(t,*args):
            R = lambda *args: abs(self.w(t,*args)) # R = |Psi|, with curried time
            numerator = laplacian(R,self.N,self.dx)
            return -(numerator(*args)) / (2*self.m*R(*args))
        return wrapper
    
    def get_Classical_Potential(self):
        return self.pot # V(x)

    def get_Hamiltonian_Bohm(self):
        # The Hamiltonian in Bohmian Mechanics is
        #
        #    H = K + V + Q
        #
        def wrapper(t,*args):
            return self.get_Kinetic_Energy()(t,*args) + self.pot(*args) + self.get_Quantum_Potential()(t,*args)
        return wrapper

    def get_Hamiltonian_Jacobi(self):
        # We can also get the Hamiltonian from the Hamilton-Jacobi equation:
        #
        #    H = -dS/dt
        #
        def wrapper(t,*args):
            S = lambda t: phase(self.w)(t,*args) # curry position
            return -partial(S,dx=self.dt)(t) # this dx is a time-step, time derivative!
        return wrapper

    def get_Hamiltonian(self):
        # in next version: do both previous with increasingly small stepsizes, only return if they agree within tolerance...
        # In the meanwhile, do the one that works most accurately...
        return self.get_Hamiltonian_Bohm()

    def get_Lagrangian(self):
        # The Lagrangian in Bohmian Mechanics is
        #
        #    L = K - V - Q
        #
        def wrapper(t,*args):
            return self.get_Kinetic_Energy()(t,*args) - self.pot(*args) - self.get_Quantum_Potential()(t,*args)
        return wrapper
            
    def get_Probability_Flux_Density(self):
        # j = rho*v = rho*(p/m)
        def wrapper(t,*args):
            return [abs(self.w(t,*args))**2 * p / self.m for p in self.get_momentum()(t,*args)]
        return wrapper
    
    def get_Force(self):
        # Uses V(x) and Q(t,x) to calculate
        #
        #    dp/dt = - dH/dx = -d(V+Q)/dx
        #
        def wrapper(t,*args):
            # curry out the time dependence
            def Q(*args):
                return self.get_Quantum_Potential()(t,*args)
            # now take the derivatives
            return -(partial(self.pot,dx=self.dx)(*args) + partial(Q,dx=self.dx)(*args))
        return wrapper
    
    # ------------------------------------------------------------------------------------
    
    # We now put together a few bits and pieces to construct trajectories from initial conditions.
    
    def __RK4_step__(self,x,t,dt=1.0):
        
        # get the Runge-Kutta coefficients from the velocities of a Schrodinger object
        
        # we only want the [0] component because we are working in 1D. In a future version
        # this will accept N dimensions and this will be self.get_mom()(t,x) without []'s
        
        k1 = self.get_momentum()(t,x)[0]
        k2 = self.get_momentum()(t+dt/2,x+dt*k1/2)[0]
        k3 = self.get_momentum()(t+dt/2,x+dt*k2/2)[0]
        k4 = self.get_momentum()(t+dt,x+dt*k3)[0]
        
        # return the next position
        return x + dt * (k1 + 2*k2 + 2*k3 + k4)/6.0 / self.m # future version: vector sum in N dimensions
    
    def __RK4_step_in_N_dimensions__(self,x,t,dt=1.0): # x a vector
        
        # get the Runge-Kutta coefficients from the velocities of a Schrodinger object
        
        # we only want the [0] component because we are working in 1D. In a future version
        # this will accept N dimensions and this will be self.get_mom()(t,x) without []'s
        
        k1 = self.get_momentum()(t,*x)
        k2 = self.get_momentum()(t+dt/2,*[i+dt*j/2.0 for i,j in zip(x,k1)])
        k3 = self.get_momentum()(t+dt/2,*[i+dt*j/2.0 for i,j in zip(x,k2)])
        k4 = self.get_momentum()(t+dt,*[i+dt*j for i,j in zip(x,k3)])
        
        # k's are vectors
        
        return map(lambda xi,a,b,c,d: xi + dt*(a+2*b+2*c+d)/6.0/self.m , x, k1, k2, k3, k4)
    
    # With Richardson Extrapolation: R(h) = A(h/2) + [A(h/2) - A(h)] / [ 2^n - 1 ]
    # In RK4 we have n = 4 (accumulated error)
    
    def __Richardson_step__(self,x,t,dt=1.0):
        x1 = self.__RK4_step__(x,t,dt)
        dt2 = dt / 2.0
        x2 = self.__RK4_step__(self.__RK4_step__(x,t,dt2), t+dt2, dt2)
        return (16*x2 - x1)/15.0 # future version: vector difference and scaling, in N dimensions
    
    def __Richardson_step_in_N_dimensions__(self,x,t,dt=1.0):
        x1 = self.__RK4_step_in_N_dimensions__(x,t,dt)
        dt2 = dt / 2.0
        x2 = self.__RK4_step_in_N_dimensions__(self.__RK4_step_in_N_dimensions__(x,t,dt2), t+dt2, dt2)
        return [(16*i - j)/15.0 for i,j in zip(x2,x1)]
    
    def get_Trajectory(self,x_ic,t_ic,t_final,timesteps):
        # Computes trajectory from initial (x,t) to the position
        # at time t_final with `timesteps' intermediate timesteps.
        
        dt = (t_final-t_ic)/timesteps #stepsize
        X=x_ic
        T=t_ic
        traj = [X]
        
        for step in range(timesteps): # time propagation
                X = self.__Richardson_step__(X,T,dt)
                T += dt
                traj.append(X)
        
        return traj

    def get_Trajectory_in_N_dimensions(self,x_ic,t_ic,t_final,timesteps):
        # Computes trajectory from initial (x,t) to the position
        # at time t_final with `timesteps' intermediate timesteps.
        
        dt = (t_final-t_ic)/timesteps #stepsize
        X=x_ic
        T=t_ic
        traj = [X]
        
        for step in range(timesteps): # time propagation
            X = self.__Richardson_step_in_N_dimensions__(X,T,dt)
            T += dt
            traj.append(X)
        
        # Current implementation is a loop to build the list. We should check
        # whether a list comprehension would be faster (if this is too slow).
        
        return traj
    
    def get_Times(self,t_ic,t_final,timesteps):
        # get_Times(0,5,10) => [0.0, 0.5, 1.0, 1.5, ... , 4.0, 4.5, 5.0]
        return [i/float(timesteps)*(t_final-t_ic)+t_ic for i in range(timesteps+1)]

    def __get_Separation__(self,X_1,T_1,X_2,T_2,Duration,timesteps):
        a = self.get_Trajectory(X_1,T_1,T_1+Duration,timesteps)
        b = self.get_Trajectory(X_2,T_2,T_2+Duration,timesteps)
        return map(lambda x,y: abs(x-y), a, b)
    
    def get_FTLES(self,x_ic,t_ic,t_final,timesteps,initial_separation=0.01):
        # Finite Time Lyapunov Exponents for various times
        s = self.__get_Separation__(x_ic+initial_separation/2.0,t_ic,x_ic-initial_separation/2.0,t_ic,t_final-t_ic,timesteps)
        t = self.get_Times(t_ic,t_final,timesteps)
        L = [(log(s[i])-log(s[0]))/(t[i]-t[0]) for i in range(1,timesteps+1)] # generate FTLE's for various final times
        return L # a list of FTLE's between t_ic and t_final

# -------------------- Testing ---------------------

if __name__ == "__main__":
    
    print("\nFirst, test the differentiation tools:\n")
    
    TOL = 1E-7 # if tolerance changes, so should the stepsize dx in the tests
    STEP_DX = 0.01
    PT_NUM = 1000
    GRID_SCALE = 10.0

    print("Checking Differentiation is OK...") # Test 'Partial'
    
    g = lambda x : 3*x**2+2*x
    G = lambda x: x**3 + x**2
    g2 = partial(G,dx=STEP_DX)
    diff = [g(i)-g2(i)<TOL for i in [j/GRID_SCALE for j in range(PT_NUM)]]
    print(all(diff))

    print("Checking Gradient is OK...")

    f = lambda x,y,z: 3*x + 2*y + z**2
    gf = lambda x,y,z : [3,2,2*z]
    gradf = gradient(f,N=3,dx=STEP_DX)
    diff = lambda i : map(lambda x,y: x-y,  gf(i,i,i), gradf(i,i,i))
    grad  = [sum(diff(i))<TOL for i in [j/GRID_SCALE for j in range(PT_NUM)]]
    print(all(grad))

    print("Checking Divergence is OK...")

    v = lambda x,y,z: [2*x,y**2,x*z]
    dv = lambda x,y,z: 2+2*y+x
    divv = divergence(v,dx=STEP_DX)
    div = [dv(i,i,i)-divv(i,i,i)<TOL for i in [j/GRID_SCALE for j in range(PT_NUM)]]
    print(all(div))

    print("Checking Laplacian is OK... (restrictions apply)")

    f = lambda x,y,z: 3*x**2 + 2*y + z**3
    lf = lambda x,y,z: 6 + 6*z
    laplf = laplacian(f,N=3,dx=STEP_DX)
    grad  = [lf(i,i,i)-laplf(i,i,i)<TOL for i in [j/1000.0 for j in range(10000)]] # 0 < j < 10
    print(all(grad))
    
    
    
    print("\nNow, test `physical' functionals:\n")
    
    
    from cmath import exp
    eigenfn = lambda t,x : 0.71637182 * x**2 * exp(-13.7*(0+1j)*t)
    eigenstate1 = SchrodingerNDT(lambda x:0, eigenfn, 1.0, lambda x: [0], 0.0, dimensions=1, dx=0.001, dt=0.01)

    from HarmonicOscillator1D import Psi4
    eigenstate2 = SchrodingerNDT(lambda x:0, Psi4, 1.0, lambda x: [0], 0.0, dimensions=1, dx=0.001, dt=0.01)
    
    print("Checking get_momentum() is OK...")
    
    p1 = eigenstate1.get_momentum()
    print(max([abs(p1(0.526,i/10.0)[0]) for i in range(-500,600)]))
    p2 = eigenstate2.get_momentum()
    print(max([abs(p2(0.526,i/10.0)[0]) for i in range(-500,600)]))
    
    print("Checking Runge-Kutta is OK for eigenstates (no motion)...")
    
    x_0 = [i/10.0 for i in range(-10,11)]
    x = x_0
    RK4 = lambda x: eigenstate1.__RK4_step__(x,0.53728,dt=0.01)
    for count_times in range(100):
        x = map(RK4,x)
    print(max([i - j for i, j in zip(x, x_0)]) < 1e-10) # look for largest difference

# TODO:
# - Do get_Hamiltonian() right.

